#define _CRT_SECURE_NO_WARNINGS
#include <fstream>
#include <iostream>
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

#include <thread>
#include <mutex>


std::mutex gMut;
cv::Mat gFrm, gRet;
cv::Mat gMask;
//using namespace cv;
//using namespace dnn;
//using namespace std;

extern const char* class_names[] = { "background",
                              "person", "bicycle", "car", "motorcycle", "airplane", "bus",
                              "train", "truck", "boat", "traffic light", "fire hydrant",
                              "stop sign", "parking meter", "bench", "bird", "cat", "dog",
                              "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe",
                              "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
                              "skis", "snowboard", "sports ball", "kite", "baseball bat",
                              "baseball glove", "skateboard", "surfboard", "tennis racket",
                              "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl",
                              "banana", "apple", "sandwich", "orange", "broccoli", "carrot",
                              "hot dog", "pizza", "donut", "cake", "chair", "couch",
                              "potted plant", "bed", "dining table", "toilet", "tv", "laptop",
                              "mouse", "remote", "keyboard", "cell phone", "microwave", "oven",
                              "toaster", "sink", "refrigerator", "book", "clock", "vase",
                              "scissors", "teddy bear", "hair drier", "toothbrush"
};

extern const unsigned char colors[81][3] = {
   {56, 0, 255},
   {226, 255, 0},
   {0, 94, 255},
   {0, 37, 255},
   {0, 255, 94},
   {255, 226, 0},
   {0, 18, 255},
   {255, 151, 0},
   {170, 0, 255},
   {0, 255, 56},
   {255, 0, 75},
   {0, 75, 255},
   {0, 255, 169},
   {255, 0, 207},
   {75, 255, 0},
   {207, 0, 255},
   {37, 0, 255},
   {0, 207, 255},
   {94, 0, 255},
   {0, 255, 113},
   {255, 18, 0},
   {255, 0, 56},
   {18, 0, 255},
   {0, 255, 226},
   {170, 255, 0},
   {255, 0, 245},
   {151, 255, 0},
   {132, 255, 0},
   {75, 0, 255},
   {151, 0, 255},
   {0, 151, 255},
   {132, 0, 255},
   {0, 255, 245},
   {255, 132, 0},
   {226, 0, 255},
   {255, 37, 0},
   {207, 255, 0},
   {0, 255, 207},
   {94, 255, 0},
   {0, 226, 255},
   {56, 255, 0},
   {255, 94, 0},
   {255, 113, 0},
   {0, 132, 255},
   {255, 0, 132},
   {255, 170, 0},
   {255, 0, 188},
   {113, 255, 0},
   {245, 0, 255},
   {113, 0, 255},
   {255, 188, 0},
   {0, 113, 255},
   {255, 0, 0},
   {0, 56, 255},
   {255, 0, 113},
   {0, 255, 188},
   {255, 0, 94},
   {255, 0, 18},
   {18, 255, 0},
   {0, 255, 132},
   {0, 188, 255},
   {0, 245, 255},
   {0, 169, 255},
   {37, 255, 0},
   {255, 0, 151},
   {188, 0, 255},
   {0, 255, 37},
   {0, 255, 0},
   {255, 0, 170},
   {255, 0, 37},
   {255, 75, 0},
   {0, 0, 255},
   {255, 207, 0},
   {255, 0, 226},
   {255, 245, 0},
   {188, 255, 0},
   {0, 255, 18},
   {0, 255, 75},
   {0, 255, 151},
   {255, 56, 0},
   {245, 255, 0}
};


class yolact
{
public:
   yolact(float confThreshold, float nmsThreshold, const int keep_top_k = 200);
   void detect(cv::Mat& srcimg);
private:
   const int target_size = 550;
   const float MEANS[3] = { 123.68, 116.78, 103.94 };
   const float STD[3] = { 58.40, 57.12, 57.38 };
   float confidence_threshold;
   float nms_threshold;
   int keep_top_k;
   const int conv_ws[5] = { 69, 35, 18, 9, 5 };
   const int conv_hs[5] = { 69, 35, 18, 9, 5 };
   const float aspect_ratios[3] = { 1.f, 0.5f, 2.f };
   const float scales[5] = { 24.f, 48.f, 96.f, 192.f, 384.f };
   const float var[4] = { 0.1f, 0.1f, 0.2f, 0.2f };
   const int mask_h = 138;
   const int mask_w = 138;
   int num_priors;
   float* priorbox;
   cv::dnn::Net net;
   void normalize(cv::Mat& img);
   void sigmoid(cv::Mat& out, int length);
};

yolact::yolact(float confThreshold, float nmsThreshold, const int keep_top_k)
{
   cv::setNumThreads(4);
   this->confidence_threshold = confThreshold;
   this->nms_threshold = nmsThreshold;
   this->keep_top_k = keep_top_k;
   this->net = cv::dnn::readNet("D:/05_cpp/00_develop_envirment/project/yolact/config/yolact_base_54_800000.onnx");
   //this->net.setPreferableBackend(cv::dnn::DNN_BACKEND_INFERENCE_ENGINE);//openvino
   this->net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
   //this->net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
   //this->net.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL);
   this->net.setPreferableTarget(cv::dnn::DNN_TARGET_OPENCL_FP16);
   this->num_priors = 0;
   int p = 0;
   for (p = 0; p < 5; p++)
   {
      this->num_priors += this->conv_ws[p] * this->conv_hs[p] * 3;
   }
   this->priorbox = new float[4 * this->num_priors];
   ////generate priorbox
   float* pb = priorbox;
   for (p = 0; p < 5; p++)
   {
      int conv_w = this->conv_ws[p];
      int conv_h = this->conv_hs[p];

      float scale = this->scales[p];

      for (int i = 0; i < conv_h; i++)
      {
         for (int j = 0; j < conv_w; j++)
         {
            // +0.5 because priors are in center-size notation
            float cx = (j + 0.5f) / conv_w;
            float cy = (i + 0.5f) / conv_h;

            for (int k = 0; k < 3; k++)
            {
               float ar = aspect_ratios[k];

               ar = sqrt(ar);

               float w = scale * ar / this->target_size;
               float h = scale / ar / this->target_size;

               // This is for backward compatability with a bug where I made everything square by accident
               // cfg.backbone.use_square_anchors:
               h = w;
               pb[0] = cx;
               pb[1] = cy;
               pb[2] = w;
               pb[3] = h;
               pb += 4;
            }
         }
      }
   }
}

void yolact::normalize(cv::Mat& img)
{
   img.convertTo(img, CV_32F);
   int i = 0, j = 0;
   for (i = 0; i < img.rows; i++)
   {
      float* pdata = (float*)(img.data + i * img.step);
      for (j = 0; j < img.cols; j++)
      {
         pdata[0] = (pdata[0] - this->MEANS[0]) / this->STD[0];
         pdata[1] = (pdata[1] - this->MEANS[1]) / this->STD[1];
         pdata[2] = (pdata[2] - this->MEANS[2]) / this->STD[2];
         pdata += 3;
      }
   }
}

void yolact::sigmoid(cv::Mat& out, int length)
{
   float* pdata = (float*)(out.data);
   int i = 0;
   for (i = 0; i < length; i++)
   {
      pdata[i] = 1.0 / (1 + expf(-pdata[i]));
   }
}

void yolact::detect(cv::Mat& srcimg)
{
   cv::setNumThreads(4);
   int img_w = srcimg.cols;
   int img_h = srcimg.rows;
   cv::Mat img;
   resize(srcimg, img, cv::Size(this->target_size, this->target_size), cv::INTER_LINEAR);
   cvtColor(img, img, cv::COLOR_BGR2RGB);
   this->normalize(img);
   cv::Mat blob = cv::dnn::blobFromImage(img);
   this->net.setInput(blob);
   std::vector<cv::Mat> outs;
   this->net.forward(outs, this->net.getUnconnectedOutLayersNames());
   /*
   cv::Mat outDetections = outs[0];
   cv::Mat outMasks = outs[1];
   const int numDetections = outDetections.size[2];
   const int numClasses = outMasks.size[1];
   outDetections = outDetections.reshape(1, outDetections.total() / 7);
   cv::Mat foregroundMask32FC1 = cv::Mat(srcimg.rows, srcimg.cols, CV_32FC1, cv::Scalar::all(0));
   for (int i = 0; i < numDetections; ++i) {
      float score = outDetections.at<float>(i, 2);
      if (score > confidence_threshold) {
         int classId = static_cast<int>(outDetections.at<float>(i, 1));
         int left = static_cast<int>(srcimg.cols * outDetections.at<float>(i, 3));
         int top = static_cast<int>(srcimg.rows * outDetections.at<float>(i, 4));
         int right = static_cast<int>(srcimg.cols * outDetections.at<float>(i, 5));
         int bottom = static_cast<int>(srcimg.rows * outDetections.at<float>(i, 6));
         left = std::max(0, std::min(left, srcimg.cols - 1));
         top = std::max(0, std::min(top, srcimg.rows - 1));
         right = std::max(0, std::min(right, srcimg.cols - 1));
         bottom = std::max(0, std::min(bottom, srcimg.rows - 1));
         cv::Rect box = cv::Rect(left, top, right - left + 1, bottom - top + 1);
         cv::Mat objectMask(outMasks.size[2], outMasks.size[3], CV_32F, outMasks.ptr<float>(i, classId));
         
         if (class_names[classId] == "person") {
            cv::Mat allZero32FC1 = cv::Mat(srcimg.rows, srcimg.cols, CV_32FC1, cv::Scalar::all(0.0));
            cv::Mat objectMaskResized;
            cv::resize(objectMask, objectMaskResized, cv::Size(box.width, box.height), 0, 0, cv::INTER_CUBIC);
            objectMaskResized.copyTo(allZero32FC1(box));
            if (allZero32FC1.rows != srcimg.rows || allZero32FC1.cols != srcimg.cols) {
               resize(allZero32FC1, allZero32FC1, cv::Size(srcimg.cols, srcimg.rows), 0, 0, cv::INTER_CUBIC);
            }
            for (unsigned short row = 0; row < allZero32FC1.rows; row++) {
               for (unsigned short col = 0; col < allZero32FC1.cols; col++) {
                  if (allZero32FC1.at<float>(row, col) - foregroundMask32FC1.at<float>(row, col) > 0) {
                     foregroundMask32FC1.at<float>(row, col) = allZero32FC1.at<float>(row, col);
                  }
               }
            }
         }
      }
   }
   gMask = foregroundMask32FC1.clone();
   */

   /////generate proposals
   std::vector<int> classIds;
   std::vector<float> confidences;
   std::vector<cv::Rect> boxes;
   std::vector<int> maskIds;
   const int num_class = outs[1].cols;
   for (int i = 0; i < this->num_priors; i++)
   {
      cv::Mat scores = outs[1].row(i).colRange(1, num_class);
      cv::Point classIdPoint;
      double score;
      // Get the value and location of the maximum score
      minMaxLoc(scores, 0, &score, 0, &classIdPoint);
      if (score > this->confidence_threshold)
      {
         const float* loc = (float*)outs[0].data + i * 4;
         const float* pb = this->priorbox + i * 4;
         float pb_cx = pb[0];
         float pb_cy = pb[1];
         float pb_w = pb[2];
         float pb_h = pb[3];

         float bbox_cx = var[0] * loc[0] * pb_w + pb_cx;
         float bbox_cy = var[1] * loc[1] * pb_h + pb_cy;
         float bbox_w = (float)(exp(var[2] * loc[2]) * pb_w);
         float bbox_h = (float)(exp(var[3] * loc[3]) * pb_h);
         float obj_x1 = bbox_cx - bbox_w * 0.5f;
         float obj_y1 = bbox_cy - bbox_h * 0.5f;
         float obj_x2 = bbox_cx + bbox_w * 0.5f;
         float obj_y2 = bbox_cy + bbox_h * 0.5f;

         // clip
         obj_x1 = std::max(std::min(obj_x1 * img_w, (float)(img_w - 1)), 0.f);
         obj_y1 = std::max(std::min(obj_y1 * img_h, (float)(img_h - 1)), 0.f);
         obj_x2 = std::max(std::min(obj_x2 * img_w, (float)(img_w - 1)), 0.f);
         obj_y2 = std::max(std::min(obj_y2 * img_h, (float)(img_h - 1)), 0.f);
         classIds.push_back(classIdPoint.x);
         confidences.push_back(score);
         boxes.push_back(cv::Rect((int)obj_x1, (int)obj_y1, (int)(obj_x2 - obj_x1 + 1), (int)(obj_y2 - obj_y1 + 1)));
         maskIds.push_back(i);
      }
   }

   // Perform non maximum suppression to eliminate redundant overlapping boxes with
   // lower confidences
   std::vector<int> indices;
   cv::dnn::NMSBoxes(boxes, confidences, this->confidence_threshold, this->nms_threshold, indices, 1.f, this->keep_top_k);
   cv::Mat mask(this->mask_h, this->mask_w, CV_32FC1);
   mask = cv::Scalar(0.f);
   float* pm = (float*)mask.data;
   int area = this->mask_h * this->mask_w;
   cv::Mat mask2 = cv::Mat(this->mask_h, this->mask_w, CV_32FC1, cv::Scalar::all(0));
   for (size_t i = 0; i < indices.size(); ++i)
   {
      int idx = indices[i];
      if (class_names[classIds[idx] + 1] == "person") {
         cv::Rect box = boxes[idx];
         int xmax = box.x + box.width;
         int ymax = box.y + box.height;
         rectangle(srcimg, cv::Point(box.x, box.y), cv::Point(xmax, ymax), cv::Scalar(0, 0, 255), 3);
         //Get the label for the class name and its confidence
         char text[256];
         sprintf(text, "%s: %.2f %%", class_names[classIds[idx] + 1], confidences[idx] * 100);

         //Display the label at the top of the bounding box
         int baseLine;
         cv::Size labelSize = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
         int ymin = std::max(box.y, labelSize.height);
         //rectangle(frame, Point(left, top - int(1.5 * labelSize.height)), Point(left + int(1.5 * labelSize.width), top + baseLine), Scalar(0, 255, 0), FILLED);
         putText(srcimg, text, cv::Point(box.x, ymin), cv::FONT_HERSHEY_SIMPLEX, 0.50, cv::Scalar(0, 255, 0), 1);

         
         int channel = outs[2].cols;
         
         float* coeff = (float*)outs[2].data + maskIds[idx] * channel;
         
         const float* pmaskmap = (float*)outs[3].data;
         for (int j = 0; j < area; j++)
         {
            for (int p = 0; p < channel; p++)
            {
               pm[j] += pmaskmap[p] * coeff[p];
            }
            pmaskmap += channel;
         }

         this->sigmoid(mask, area);

         resize(mask, mask2, cv::Size(img_w, img_h), 0, 0, cv::INTER_CUBIC);
         
         // draw mask
         for (int y = 0; y < img_h; y++)
         {
            const float* pmask = (float*)mask2.data + y * img_w;
            uchar* p = srcimg.data + y * img_w * 3;
            for (int x = 0; x < img_w; x++)
            {
               if (pmask[x] > 0.5)
               {
                  p[0] = (uchar)(p[0] * 0.5 + colors[classIds[idx] + 1][0] * 0.5);
                  p[1] = (uchar)(p[1] * 0.5 + colors[classIds[idx] + 1][1] * 0.5);
                  p[2] = (uchar)(p[2] * 0.5 + colors[classIds[idx] + 1][2] * 0.5);
               }
               p += 3;
            }
         }
         gMask = mask2.clone();
      }
   }

  
}

#include <Windows.h>

yolact* gYolact = new yolact(0.6, 0.7);
bool segStatus = false, gStop = false;
void segment() {
   while (true && gYolact && !gStop) {
      gMut.lock();
      DWORD startTime = GetTickCount();
      gYolact->detect(gFrm);
      gRet = gFrm.clone();
      segStatus = true;
      DWORD endTime = GetTickCount();
      std::cout << "\t\t\t\t [1] inference time: " << (endTime - startTime) << " ms" << std::endl;
      gMut.unlock();
   }
}
std::thread segmentTh() {
   return std::thread(&segment);
}

cv::Mat gFrm1;
yolact* gYolact1 = new yolact(0.6, 0.7);
bool segStatus1 = false, gStop1 = false;
void segment1() {
   while (true && gYolact1 && !gStop1) {
      gMut.lock();
      DWORD startTime = GetTickCount();
      gYolact1->detect(gFrm1);
      gRet = gFrm1.clone();
      segStatus1 = true;
      DWORD endTime = GetTickCount();
      std::cout << "\t\t\t\t [2] inference2 time: " << (endTime - startTime) << " ms" << std::endl;
      gMut.unlock();
   }
}
std::thread segmentTh1() {
   return std::thread(&segment1);
}


cv::Mat convertMaskC1ToC3(cv::Mat& mask) {
   cv::Mat* matC1Vector = new cv::Mat[3];
   for (unsigned char i = 0; i < 3; i++) {
      matC1Vector[i] = mask.clone();
   }
   cv::Mat matC3;
   cv::merge(matC1Vector, 3, matC3);
   if (matC1Vector) {
      delete[] matC1Vector;
   }
   return matC3;
}


int main()
{
   cv::setNumThreads(4);
   DWORD timerStart, timerEnd, timerUse;
   unsigned short fps = 0;
   unsigned short width = 640, height = 480;
   cv::Mat backgroundImg = cv::imread("D:/05_cpp/00_develop_envirment/project/virtual_background/test_data/500_313.jpg");
   cv::resize(backgroundImg, backgroundImg, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
   //yolact yolactnet(0.5, 0.5);
   cv::VideoCapture cap(0, cv::CAP_DSHOW);
   //cv::VideoCapture cap("D:/05_cpp/00_develop_envirment/project/virtual_background/test_data/1280_720.mp4");
   cv::Mat srcimg;
   cap >> srcimg;
   cv::resize(srcimg, srcimg, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
   cv::Mat OneMat = cv::Mat(height, width, CV_32FC1, cv::Scalar::all(1));
   gMask = cv::Mat(height, width, CV_32FC1, cv::Scalar::all(0));
   gFrm = srcimg.clone();
   gFrm1 = srcimg.clone();
   gRet = srcimg.clone();

   cv::Mat foreMaskC1, backMaskC1;
   cv::Mat foreMaskC3, backMaskC3;
   cv::Mat foreMask32f, backMask32f;
   foreMaskC1 = gMask.clone();
   cv::resize(foreMaskC1, foreMaskC1, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
   cv::subtract(OneMat, foreMaskC1, backMaskC1);
   foreMaskC3 = convertMaskC1ToC3(foreMaskC1);
   backMaskC3 = convertMaskC1ToC3(backMaskC1);
   foreMaskC3.convertTo(foreMask32f, CV_32FC3);
   backMaskC3.convertTo(backMask32f, CV_32FC3);

   std::thread segTh = segmentTh();
   if (segTh.joinable()) {
      segTh.detach();
   }
   //std::thread segTh1 = segmentTh1();
   //if (segTh1.joinable()) {
   //   segTh1.detach();
   //}
   static const std::string kWinName = "yolact";
   cv::namedWindow(kWinName, cv::WINDOW_NORMAL);
   while (cap.isOpened()) {
      timerStart = GetTickCount();
      cap >> srcimg;
      if (srcimg.cols != width || srcimg.rows != height) {
         cv::resize(srcimg, srcimg, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
      }
      cv::imshow("input", srcimg);
      //yolactnet.detect(srcimg);
      if (segStatus) { 
         gFrm = srcimg.clone();
         foreMaskC1 = gMask.clone();
         cv::resize(foreMaskC1, foreMaskC1, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
         cv::subtract(OneMat, foreMaskC1, backMaskC1);
         foreMaskC3 = convertMaskC1ToC3(foreMaskC1);
         backMaskC3 = convertMaskC1ToC3(backMaskC1);
         foreMaskC3.convertTo(foreMask32f, CV_32FC3);
         backMaskC3.convertTo(backMask32f, CV_32FC3);
         segStatus = false;
      }
      if (segStatus1) {
         gFrm1 = srcimg.clone();
         foreMaskC1 = gMask.clone();
         cv::resize(foreMaskC1, foreMaskC1, cv::Size(width, height), 0, 0, cv::INTER_CUBIC);
         cv::subtract(OneMat, foreMaskC1, backMaskC1);
         foreMaskC3 = convertMaskC1ToC3(foreMaskC1);
         backMaskC3 = convertMaskC1ToC3(backMaskC1);
         foreMaskC3.convertTo(foreMask32f, CV_32FC3);
         backMaskC3.convertTo(backMask32f, CV_32FC3);
         segStatus1 = false;
      }
      cv::Mat srcimg32f, backgroundimg32f, foreret32f, backret32f;
      srcimg.convertTo(srcimg32f, CV_32FC3);
      backgroundImg.convertTo(backgroundimg32f, CV_32FC3);
      foreret32f = srcimg32f.mul(foreMask32f);
      backret32f = backgroundimg32f.mul(backMask32f);
      cv::Mat ret32f, ret8u;
      cv::add(foreret32f, backret32f, ret32f);
      ret32f.convertTo(ret8u, CV_8UC3);
      imshow(kWinName, ret8u);
      srcimg = gRet.clone();
      cv::imshow("segment ret", srcimg);

      timerEnd = GetTickCount();
      timerUse = timerEnd - timerStart;
      fps = (timerUse > 0) ? ((unsigned short)(1000 / timerUse)) : 999;
      std::cout << "time: " << timerUse << " ms\tfps: " << fps << std::endl;
      unsigned short waitTime = (timerUse >= 66) ? 1 : (66 - timerUse);
      if (cv::waitKey(waitTime) == 27) {
         gStop = true;
         gStop1 = true;
         Sleep(3000);
         break;
      }
   }
   cv::destroyAllWindows();
   cap.release();
   if (gYolact) {
      delete gYolact;
   }
   if (gYolact1) {
      delete gYolact1;
   }
   system("pause");
   return 0;
}

